import os
import argparse


def get_parser():
    parser = argparse.ArgumentParser(description="Training config for FontDiffuser.")
    ################# Experience #################
    parser.add_argument("--seed", type=int, default=123, help="A seed for reproducible training.")
    parser.add_argument("--experience_name", type=str, default="fontdiffuer_training")
    parser.add_argument("--data_root", type=str, default=None,
                        help="The font dataset root path.", )
    parser.add_argument("--output_dir", type=str, default=None,
                        help="The output directory where the model predictions and checkpoints will be written.")
    parser.add_argument("--report_to", type=str, default="tensorboard")
    parser.add_argument("--logging_dir", type=str, default="logs",
                        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
                              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."))

    # Model
    parser.add_argument("--resolution", type=int, default=96,
                        help="The resolution for input images, all the images in the train/validation \
                            dataset will be resized to this.")
    parser.add_argument("--unet_channels", type=tuple, default=(64, 128, 256, 512),
                        help="The channels of the UNet.")
    parser.add_argument("--style_image_size", type=int, default=96, help="The size of style images.")
    parser.add_argument("--content_image_size", type=int, default=96, help="The size of content images.")
    parser.add_argument("--content_encoder_downsample_size", type=int, default=3,
                        help="The downsample size of the content encoder.\
                        指定内容编码器的下采样尺寸，默认值为3。这意味着输入的特征图会被下采样到1/3的大小，控制模型对内容图像的细节提取能力。")
    parser.add_argument("--channel_attn", type=bool, default=True, help="Whether to use the se attention.\
    是否使用通道注意力机制（SE注意力机制），默认为True。SE注意力机制有助于增强模型对重要特征的关注。", )
    parser.add_argument("--content_start_channel", type=int, default=64,
                        help="The channels of the fisrt layer output of content encoder.", )
    parser.add_argument("--style_start_channel", type=int, default=64,
                        help="The channels of the fisrt layer output of content encoder.", )

    # Training
    parser.add_argument("--phase_2", action="store_true", help="Training in phase 2 using SCR module.")
    parser.add_argument("--phase_1_ckpt_dir", type=str, default=None, help="The trained ckpt directory during phase 1.")
    ## SCR
    parser.add_argument("--temperature", type=float, default=0.07)
    parser.add_argument("--mode", type=str, default="refinement")
    parser.add_argument("--scr_image_size", type=int, default=96)
    parser.add_argument("--scr_ckpt_path", type=str, default=None)
    parser.add_argument("--num_neg", type=int, default=16, help="Number of negative samples.\
    指定负样本的数量。负样本通常用于对比学习或类似的任务中，帮助模型学习区分正负样本的特征。")
    parser.add_argument("--nce_layers", type=str, default='0,1,2,3',
                        help='在对比学习中使用的层。NCE通常指的是噪声对比估计（Noise Contrastive Estimation）')
    parser.add_argument("--sc_coefficient", type=float, default=0.01)
    ## train batch size
    parser.add_argument("--train_batch_size", type=int, default=4,
                        help="Batch size (per device) for the training dataloader.")
    ## loss coefficient
    parser.add_argument("--perceptual_coefficient", type=float, default=0.01)
    parser.add_argument("--offset_coefficient", type=float, default=0.5)
    ## step
    parser.add_argument("--max_train_steps", type=int, default=440000,
                        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.", )
    parser.add_argument("--ckpt_interval", type=int, default=40000, help="The step begin to validate.")
    parser.add_argument("--gradient_accumulation_steps", type=int, default=1,
                        help="Number of updates steps to accumulate before performing a backward/update pass.", )
    parser.add_argument("--log_interval", type=int, default=100, help="The log interval of training.")
    ## learning rate
    parser.add_argument("--learning_rate", type=float, default=1e-4,
                        help="Initial learning rate (after the potential warmup period) to use.")
    parser.add_argument("--scale_lr", action="store_true", default=False,
                        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.\
                        如果启用 (--scale_lr)，学习率会根据这些因素自动调整，适应分布式训练的情况。否则，学习率将保持为固定值，不进行缩放。\
                        这在多GPU或分布式训练场景下非常有用，因为批量大小和计算资源的增加通常需要相应地调整学习率来保持训练稳定性。")
    parser.add_argument("--lr_scheduler", type=str, default="linear",
                        help="The scheduler type to use. Choose between 'linear', 'cosine', \
                            'cosine_with_restarts', 'polynomial', 'constant', 'constant_with_warmup'")
    parser.add_argument("--lr_warmup_steps", type=int, default=10000,
                        help="Number of steps for the warmup in the lr scheduler.")
    ## classifier-free
    parser.add_argument("--drop_prob", type=float, default=0.1, help="The uncondition training drop out probability.")
    ## scheduler
    parser.add_argument("--beta_scheduler", type=str, default="scaled_linear", help="The beta scheduler for DDPM.")
    ## optimizer
    parser.add_argument("--adam_beta1", type=float, default=0.9, help="The beta1 parameter for the Adam optimizer.")
    parser.add_argument("--adam_beta2", type=float, default=0.999, help="The beta2 parameter for the Adam optimizer.")
    parser.add_argument("--adam_weight_decay", type=float, default=1e-2, help="Weight decay to use.")
    parser.add_argument("--adam_epsilon", type=float, default=1e-08, help="Epsilon value for the Adam optimizer")
    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
    """
    1. no: 不使用混合精度，全部计算都以 32 位浮点数（fp32）进行。此模式最稳定，但效率相对较低。
    2. fp16: 使用半精度浮点数（float16 precision，也称 fp16）进行计算。fp16 占用较少的内存，且可以加速训练，但需要硬件支持（如 Nvidia Volta、Ampere GPU）来处理精度可能导致的问题。
    3. bf16: 使用 bfloat16（半精度浮点数的一种变体）。相比 fp16，bf16 保留更多的指数位，计算稳定性更高，同时具备类似的加速效果。bf16 需要 PyTorch 版本 >= 1.10，并且通常需要 Ampere GPU。
    """
    parser.add_argument("--mixed_precision", type=str, default="no", choices=["no", "fp16", "bf16"],
                        help="Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires \
                            PyTorch >= 1.10. and an Nvidia Ampere GPU.\
                            用于指定训练过程中使用的混合精度模式。混合精度训练可以加速深度学习模型的训练过程")

    # Sampling
    parser.add_argument("--algorithm_type", type=str, default="dpmsolver++", help="Algorithm for sampleing.")
    parser.add_argument("--guidance_type", type=str, default="classifier-free", help="Guidance type of sampling.")
    parser.add_argument("--guidance_scale", type=float, default=7.5, help="Guidance scale of the classifier-free mode.")
    parser.add_argument("--num_inference_steps", type=int, default=20, help="Sampling step.")
    parser.add_argument("--model_type", type=str, default="noise", help="model_type for sampling.")
    parser.add_argument("--order", type=int, default=2, help="The order of the dpmsolver.")
    parser.add_argument("--skip_type", type=str, default="time_uniform", help="Skip type of dpmsolver.")
    parser.add_argument("--method", type=str, default="multistep", help="Multistep of dpmsolver.")
    parser.add_argument("--correcting_x0_fn", type=str, default=None, help="correcting_x0_fn of dpmsolver.")
    parser.add_argument("--t_start", type=str, default=None, help="t_start of dpmsolver.")
    parser.add_argument("--t_end", type=str, default=None, help="t_end of dpmsolver.")

    """
    local_rank 参数在分布式训练中用于指定每个 GPU 或计算节点的标识。它的主要作用是在多 GPU 或分布式训练时，帮助框架区分每个设备的任务。
    详细作用：
    分布式训练：在大规模模型训练中，通常会使用多 GPU 或多机器来并行处理数据。local_rank 参数用于为每个 GPU 分配一个唯一的标识，以确保在多 GPU 环境下，任务和计算可以正确分配给每个设备。
    通信同步：当进行梯度计算时，多个 GPU 需要彼此同步参数。local_rank 帮助每个节点知道自己在整个集群中的位置，从而正确参与数据的通信、计算和梯度同步。
    默认值为 -1：这表示单 GPU 训练或未启用分布式模式。在这种情况下，local_rank 并不会实际起作用。
    示例：
    当使用 torch.distributed.launch 启动分布式训练时，local_rank 会被自动设置为每个进程对应的 GPU 号。
    python -m torch.distributed.launch --nproc_per_node=4 train.py --local_rank=0
    每个进程会接收不同的 local_rank 值（如 0, 1, 2, 3），确保任务正确分配到不同的 GPU。
    """
    parser.add_argument("--local_rank", type=int, default=-1, help="For distributed training: local_rank")

    return parser
